# coding:utf-8
import tensorflow as tf
import cv2
import os
from keras.preprocessing import image
from imageio import imread
from matplotlib import pyplot as plt
import numpy as np
import time

path=r"C:\Users\Administrator\Desktop\frsop.pb"
PATH_TO_TEST_IMAGES_DIR = r'C:\Users\Administrator\Desktop\screw-ng1\0'
TEST_IMAGE_PATHS = [os.path.join(PATH_TO_TEST_IMAGES_DIR, i) for i in os.listdir(PATH_TO_TEST_IMAGES_DIR)]

detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef() # 获取序列化后的图
    # 载入pb文件
    with tf.gfile.GFile(path,'rb') as fid:
        serialized_graph=fid.read()
        # 解析内容
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def,name='')
# REW:打印节点
# print(detection_graph.get_operations())
# exit(0)
input_1 = detection_graph.get_tensor_by_name('input_1:0')  # 获取名字
oupu = detection_graph.get_tensor_by_name('output_1:0')
input_images = []
orig_images =[]
confidence_threshold = 0.5
image_size = [200,200]
classes = ["BG", "screw"]
print(len(classes))
plt.figure(figsize=(20, 12))
with tf.Session(graph=detection_graph) as sess:
    for image_path in TEST_IMAGE_PATHS[:30]:
        orimg = imread(image_path)
        img = image.load_img(image_path, target_size=(image_size[0], image_size[1]))
        img = image.img_to_array(img)
        # input_images.append(img)
        img = np.expand_dims(img,axis=0)
    # input_images = np.array(input_images)[0]
    # input_images = np.expand_dims(input_images,axis=0)
        # input_1 = detection_graph.get_tensor_by_name('conv2d_1_input:0')
        # t = detection_graph.get_tensor_by_name('dropout_1/keras_learning_phase:0')

    # prediction = sess.run([oupu], feed_dict={input_1: gray})
        y_pred = sess.run(oupu, feed_dict={input_1: img})
        print(y_pred.shape)

        y_pred_thresh = y_pred[[y_pred[:, 1] > confidence_threshold]]

        np.set_printoptions(precision=2, suppress=True, linewidth=90)
        print("Predicted boxes:\n")
        print('   class   conf xmin   ymin   xmax   ymax')
        # for i,j in enumerate(y_pred_thresh):
        #     print(i,"张图片",j)
        #     print("\n\n")

        colors = plt.cm.hsv(np.linspace(0, 1, 81)).tolist()

        plt.imshow(orimg)
        current_axis = plt.gca()

        for box in y_pred_thresh:
            # Transform the predicted bounding boxes for the 300x300 image to the original image dimensions.
            xmin = box[2] * orimg.shape[1] / image_size[0]  # x*ori_shape / 300 = y
            ymin = box[3] * orimg.shape[0] / image_size[1]
            xmax = box[4] * orimg.shape[1] / image_size[1]
            ymax = box[5] * orimg.shape[0] / image_size[0]
            color = colors[int(box[0])]
            label = '{}: {:.2f}'.format(classes[int(box[0])], box[1])
            current_axis.add_patch(
                plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False, linewidth=2))
            current_axis.text(xmin, ymin, label, size='x-large', color='white', bbox={'facecolor': color, 'alpha': 1.0})
        plt.show()
        time.sleep(1)
        current_axis.clear()
